﻿/**
 * 基于 ZeroMQ 的 RPC
 *
 * 2021.12.24 薄云峰 Yunfeng.Bo@qq.com
 *
 * 序列化使用 msgpack：
 *   - 不使用 boost.serializer 或自定义的序列化等 C++
 * 特有的序列化库，是为了兼容更多语言；
 *   - 不使用基于 json 的序列化库，是考虑性能问题；
 *   - 不使用 protobuf 等，是因为考虑到 msgpack 开发更简单一点，不需要 proto
 * 等接口定义文件；
 */

// TODO 目前只实现了单服务线程
// TODO 目前只实现了基于REQ/REP模式的调用
// TODO 目前只提供了基于TCP方式的通信

#pragma once

#include <functional>
#include <iostream>
#include <map>
#include <thread>
#include <tuple>

#include "serializer.hpp"
#include <zmq.hpp>

enum class rpc_role {
  RPC_CLIENT, // client 端
  RPC_SERVER  // server 端
};

enum class rpc_errcode {
  RPC_ERR_SUCCESS = 0,          // 调用成功
  RPC_ERR_FUNCTION_NOT_REGUIST, // 调用的方法未注册
  RPC_ERR_RECV_TIMEOUT,         // 调用超时
  RPC_ERR_ARGS_COUNT_ERR, // 参数不足（包括没指定方法名，也认为“参数不足”）
  RPC_ERR_ARGS_TYPE_ERR, // 参数类型错误，无法正确转换为注册方法的参数类型
  RPC_ERR_UNKNOWN, // 发生未知类型错误
};

// 内部管理用，注册为 RPC 服务的函数
// 第一个参数用于返回调用结果
// 第二个参数是调用的参数
typedef std::function<void(msgpack_serializer *, msgpack_serializer &)>
    rpc_func;

template <typename T> struct type_xx { typedef T type; };

template <> struct type_xx<void> { typedef int8_t type; };

// 打包函数调用参数，辅助
template <typename Tuple, std::size_t... Index>
void package_args_impl(msgpack_serializer &sr, const Tuple &t,
                       std::index_sequence<Index...>) {
  std::initializer_list<int>{(sr << std::get<Index>(t), 0)...};
}

// 打包函数调用参数
template <typename... Args>
void package_args(msgpack_serializer &sr, const std::tuple<Args...> &t) {
  package_args_impl(sr, t, std::index_sequence_for<Args...>{});
}

// 用 Tuple 做参数调用函数模板类，辅助
template <typename Function, typename Tuple, std::size_t... Index>
decltype(auto) invoke_impl(Function &&func, Tuple &&t,
                           std::index_sequence<Index...>) {
  return func(std::get<Index>(std::forward<Tuple>(t))...);
}

// 用 Tuple 做参数调用函数模板类
template <typename Function, typename Tuple>
decltype(auto) invoke(Function &&func, Tuple &&t) {
  constexpr auto size =
      std::tuple_size<typename std::decay<Tuple>::type>::value;
  return invoke_impl(std::forward<Function>(func), std::forward<Tuple>(t),
                     std::make_index_sequence<size>{});
}

// 调用服务方法
template <typename R, typename F, typename ArgsTuple>
typename std::enable_if<!std::is_same<R, void>::value,
                        typename type_xx<R>::type>::type
call_helper(F f, ArgsTuple args) {
  return invoke(f, args);
}

// 用于封装远程调用的结果
template <typename T> class response_t {
public:
  response_t() : m_code((int)rpc_errcode::RPC_ERR_SUCCESS) {
    m_message.clear();
  }
  int code() const { return m_code; }
  std::string message() const { return m_message; }
  T value() const { return m_value; }

  void set_value(const T &val) { m_value = val; }
  void set_code(int code) { m_code = (int)code; }
  void set_message(const std::string &msg) { m_message = msg; }

  friend msgpack_serializer &operator>>(msgpack_serializer &in,
                                        response_t<T> &d) {
    in >> d.m_code >> d.m_message;

    // 只有调用成功，才会包含有意义的返回值
    if (d.m_code == (int)rpc_errcode::RPC_ERR_SUCCESS)
      in >> d.m_value;

    return in;
  }

  friend msgpack_serializer &operator<<(msgpack_serializer &out,
                                        response_t<T> &d) {
    out << d.m_code << d.m_message << d.m_value;
    return out;
  }

private:
  int m_code;
  std::string m_message;
  T m_value;
};

// 基于 ZeroMQ 和 MsgPack 的 RPC，同时提供客户端和服务端
class zmq_rpc {
public:
  zmq_rpc()
      : m_ctx(1), m_role(rpc_role::RPC_SERVER), m_stop_running(false),
        m_service_thread(nullptr) {}

  ~zmq_rpc() {
    stop();

    m_socket->close();
    m_ctx.close();
  }

  /**
   * @brief 设置为 client 端
   * @param addr 服务端地址（tcp://127.0.0.1:[port] | inproc://somename）
   */
  void as_client(const std::string &addr) {
    m_role = rpc_role::RPC_CLIENT;
    m_socket =
        std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t *)>>(
            new zmq::socket_t(m_ctx, ZMQ_REQ), [](zmq::socket_t *sock) {
              sock->close();
              delete sock;
              sock = nullptr;
            });
    m_socket->connect(addr);
  }

  /**
   * @brief 设置为 server 端
   * @param addr 服务绑定地址（tcp://*:[port] | inproc://somename）
   */
  void as_server(const std::string &addr) {
    m_role = rpc_role::RPC_SERVER;
    m_socket =
        std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t *)>>(
            new zmq::socket_t(m_ctx, ZMQ_REP), [](zmq::socket_t *sock) {
              sock->close();
              delete sock;
              sock = nullptr;
            });
    LOG_INFO << "ZeroMQ RPC server start at : " << addr << std::endl;
    m_socket->bind(addr);
  }

  /**
   * @brief 设置超时
   * @param ms 超时时长（毫秒）
   */
  void set_timeout(int ms) {
    if (m_role == rpc_role::RPC_CLIENT)
      m_socket->set(zmq::sockopt::rcvtimeo, ms);
  }

  /**
   * @brief client 使用，调用 RPC 服务
   * @tparam R 返回值类型
   * @tparam ...Args
   * @param name RPC 方法名
   * @param ...args
   * @return response_t 类型，包含返回的 code、message 和 value
   */
  template <typename R, typename... Args>
  response_t<R> call(const std::string &name, Args... args) {
    using args_type = std::tuple<typename std::decay<Args>::type...>;
    args_type as = std::make_tuple(args...);
    msgpack_serializer sr;
    sr << name;
    package_args(sr, as);
    return net_call<R>(sr);
  }

  /**
   * @brief client 使用，网络透传代理：不对数据和结果进行解包，只是透传
   * @param data RPC 调用输入数据（包换方法名和参数，不过这里不管）
   * @param length RPC 调用输入数据长度
   * @return RPC 调用返回的二进制的数据，放在 std::string
   */
  std::string net_proxy(const char *data, size_t length) {
    zmq::message_t request(length + 1);
    memcpy(request.data(), data, length);
    send(request);

    zmq::message_t reply;
    if (!recv(reply)) {
      response_t<std::string> val;
      val.set_code((int)rpc_errcode::RPC_ERR_RECV_TIMEOUT);
      val.set_message("recv timeout");

      msgpack_serializer res;
      res << val;
      return std::string(res.data(), res.size());
    }

    std::string str = std::string((const char *)reply.data(), reply.size());
    return str;
  }

  /**
   * @brief server 使用，注册一个变通函数成为 RPC 服务方法
   * @tparam F
   * @param name client 调用时使用的方法名
   * @param func 服务响应方法
   */
  template <typename F> void regist(const std::string &name, F func) {
    m_map_funcs[name] = std::bind(&zmq_rpc::callproxy<F>, this, func,
                                  std::placeholders::_1, std::placeholders::_2);
  }

  /**
   * @brief server 使用，注册一个类成员函数成为 RPC 服务方法
   * @tparam F
   * @tparam S
   * @param name client 调用时使用的方法名
   * @param func 服务响应方法，类成员函数 &class::function
   * @param s 提供响应的类实例 &class_object
   */
  template <typename F, typename S>
  void regist(const std::string &name, F func, S *s) {
    m_map_funcs[name] = std::bind(&zmq_rpc::callproxy<F, S>, this, func, s,
                                  std::placeholders::_1, std::placeholders::_2);
  }

  /**
   * @brief server 使用，一切准备就绪后，调用 run，启动 RPC 服务
   */
  void run() {
    if (m_service_thread)
      stop();

    m_service_thread = std::make_shared<std::thread>([&]() {
      while (!m_stop_running) {
        // 接收数据
        zmq::message_t data;
        recv(data);

        // RPC 方法名
        std::string funcname;

        // 封装返回的数据
        msgpack_serializer::ptr rptr(nullptr);

        try {
          // 准备序列化
          msgpack_serializer sr((char *)data.data(), data.size());

          // 先取出调用的方法名称
          sr >> funcname;

          // 尝试调用名称指定的服务方法，并获得返回的 response
          rptr = call_(funcname, sr);

          LOG_INFO << "ZMQ_RPC Call finished, func: " << funcname;
        } catch (error_args_count er) {
          // 参数错误（或者甚至没有方法名）
          LOG_WARNING << "ZMQ_RPC Call args not enough, func: " << funcname
                      << ", error mesgae: " << er.what();
          rptr = std::make_shared<msgpack_serializer>();
          *rptr.get() << static_cast<int>(rpc_errcode::RPC_ERR_ARGS_COUNT_ERR);
          *rptr.get() << std::string(er.what());
        } catch (error_args_type er) {
          // 参数类型错误（无法转换）
          LOG_WARNING << "ZMQ_RPC Call args type error, func: " << funcname
                      << ", error mesgae: " << er.what();
          rptr = std::make_shared<msgpack_serializer>();
          *rptr.get() << static_cast<int>(rpc_errcode::RPC_ERR_ARGS_TYPE_ERR);
          *rptr.get() << std::string(er.what());
        } catch (...) {
          // 其它未知错误
          LOG_WARNING << "ZMQ_RPC Call unknown error, func: " << funcname;
          rptr = std::make_shared<msgpack_serializer>();
          *rptr.get() << static_cast<int>(rpc_errcode::RPC_ERR_UNKNOWN);
          *rptr.get() << std::string("Unknown error!");
        }

        // 发送 response，如果返回数据为空，则不返回，客户端需要正确处理超时
        if (rptr) {
          zmq::message_t response(rptr->data(), rptr->size());
          send(response);
        } else {
          LOG_ERROR << "ZMQ_RPC Call, but somthing error, no response!";
        }
      }
    });
  }

  /**
   * @brief 停止服务
   */
  void stop() {
    if (!m_service_thread)
      return;

    m_stop_running = true;
    if (m_service_thread->joinable())
      m_service_thread->join();

    m_service_thread.reset();
  }

private:
  /**
   * @brief server 端，根据方法名称，查找并调用注册的服务方法
   * @param name 被调用的RPC方法名
   * @param sr_args 序列化的参数
   * @return 返回打包好的（序列化后）的 response
   */
  msgpack_serializer::ptr call_(const std::string &name,
                                msgpack_serializer &sr_args) {
    msgpack_serializer::ptr sp = std::make_shared<msgpack_serializer>();

    if (m_map_funcs.find(name) == m_map_funcs.end()) {
      LOG_WARNING << "RPC call, but function not regist, name = " << name
                  << std::endl;
      *sp.get() << static_cast<int>(rpc_errcode::RPC_ERR_FUNCTION_NOT_REGUIST);
      *sp.get() << std::string("function not regist : " + name);
      return sp;
    }

    auto func = m_map_funcs[name];

    // 调用服务方法
    func(sp.get(), sr_args);

    return sp;
  }

  template <typename F>
  void callproxy(F func, msgpack_serializer *pr, msgpack_serializer &sr_args) {
    callproxy_(func, pr, sr_args);
  }

  template <typename F, typename S>
  void callproxy(F func, S *s, msgpack_serializer *pr,
                 msgpack_serializer &sr_args) {
    callproxy_(func, s, pr, sr_args);
  }

  template <typename R, typename... Args>
  void callproxy_(R (*func)(Args...), msgpack_serializer *pr,
                  msgpack_serializer &sr_args) {
    callproxy_(std::function<R(Args...)>(func), pr, sr_args);
  }

  template <typename R, typename C, typename S, typename... Args>
  void callproxy_(R (C::*func)(Args...), S *s, msgpack_serializer *pr,
                  msgpack_serializer &sr_args) {
    using args_type = std::tuple<typename std::decay<Args>::type...>;

    args_type as =
        sr_args.get_tuple<args_type>(std::index_sequence_for<Args...>{});

    auto ff = [=](Args... args) -> R { return (s->*func)(args...); };
    typename type_xx<R>::type r = call_helper<R>(ff, as);

    response_t<R> response;
    response.set_code((int)rpc_errcode::RPC_ERR_SUCCESS);
    response.set_message("success");
    response.set_value(r);
    (*pr) << response;
  }

  template <typename R, typename... Args>
  void callproxy_(std::function<R(Args... args)> func, msgpack_serializer *pr,
                  msgpack_serializer &sr_args) {
    using args_type = std::tuple<typename std::decay<Args>::type...>;

    args_type as =
        sr_args.get_tuple<args_type>(std::index_sequence_for<Args...>{});

    typename type_xx<R>::type r = call_helper<R>(func, as);

    response_t<R> response;
    response.set_code((int)rpc_errcode::RPC_ERR_SUCCESS);
    response.set_message("success");
    response.set_value(r);
    (*pr) << response;
  }

  /**
   * @brief client 端，发送调用请求，然后接收服务端响应
   * @tparam R 返回值类型
   * @param sr 打包好的调用请求，包含方法名称和参数
   * @return 接收到响应后，解包到 response_t
   */
  template <typename R> response_t<R> net_call(msgpack_serializer &sr) {
    zmq::message_t request(sr.size() + 1);
    memcpy(request.data(), sr.data(), sr.size());
    send(request);
    zmq::message_t reply;
    response_t<R> val;
    if (!recv(reply)) {
      val.set_code((int)rpc_errcode::RPC_ERR_RECV_TIMEOUT);
      val.set_message("recv timeout");
      return val;
    }
    msgpack_serializer res((const char *)reply.data(), reply.size());
    res >> val;
    return val;
  }

  /**
   * @brief 调用 ZeroMQ，发送 RPC 调用请求
   * @param data 打包后的 RPC 调用请求，用于发送
   * @return 返回是否发送成功
   */
  bool send(zmq::message_t &data) {
    zmq::send_result_t result = m_socket->send(data, zmq::send_flags::none);
    return result.has_value();
  }

  /**
   * @brief 调用 ZeroMQ，接收 RPC 调用返回数据
   * @param data 接收到的数据
   * @return 接收是否成功
   */
  bool recv(zmq::message_t &data) {
    zmq::recv_result_t result = m_socket->recv(data, zmq::recv_flags::none);
    return result.has_value();
  }

private:
  std::map<std::string, rpc_func> m_map_funcs;
  zmq::context_t m_ctx;
  std::unique_ptr<zmq::socket_t, std::function<void(zmq::socket_t *)>> m_socket;
  rpc_role m_role;
  bool m_stop_running;
  std::shared_ptr<std::thread> m_service_thread;
};